package org.kiwi.data;

import com.alibaba.druid.pool.DruidDataSource;
import org.apache.commons.lang3.StringUtils;
import org.kiwi.data.lb.LoadBalanceStrategy;
import org.kiwi.data.spi.DataSourceWarningProvider;
import org.kiwi.data.util.DataSourceUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.util.Assert;

import javax.annotation.PreDestroy;
import javax.management.MBeanServer;
import javax.management.ObjectName;
import javax.sql.DataSource;
import java.lang.management.ManagementFactory;
import java.lang.reflect.Field;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

/**
 * Created by jack on 16/10/19.
 */
public class RoutingDataSource extends AbstractRoutingDataSource implements RoutingDataSourceMBean {

    private static final Logger LOGGER = LoggerFactory.getLogger(RoutingDataSource.class);

    private Map<Object, DataSource> resolvedDataSources;

    @Override
    public void afterPropertiesSet() {
        super.afterPropertiesSet();

        try {
            Field field = this.getClass().getSuperclass().getDeclaredField("resolvedDataSources");
            field.setAccessible(true);
            this.resolvedDataSources = (Map<Object, DataSource>) field.get(this);
        } catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }

        Assert.notNull(this.resolvedDataSources, "resolvedDataSources is required");

        registerRoutingDataSourceMBean();
    }

    private void registerRoutingDataSourceMBean() {
        synchronized (this) {
            try {
                String mbeanName = this.getClass().getSimpleName();

                MBeanServer mbeanServer = ManagementFactory.getPlatformMBeanServer();
                ObjectName objectName = new ObjectName("org.kiwi.data:type=" + mbeanName);

                if (!mbeanServer.isRegistered(objectName)) {
                    mbeanServer.registerMBean(this, objectName);
                    if (LOGGER.isDebugEnabled()) {
                        LOGGER.debug("org.kiwi.data:type=" + mbeanName + " registered successfully");
                    }
                }
            } catch (Exception e) {
                LOGGER.error(e.getMessage());
            }
        }
    }

    @Override
    protected Object determineCurrentLookupKey() {
        String currentLookupKey = ReadWriteSeparationHolder.getReadWriteDatasourceName();

        if (StringUtils.isBlank(currentLookupKey)) {
            //本次请求未配置读写分离,不会被读写分离拦截器拦截,故其中的分库路由逻辑不会被执行
            //执行分库逻辑并选举数据源
            String groupName = DbShardingStrategyHolder.getDbGroupName();
            if (StringUtils.isNotBlank(groupName)) {

                //只开启了分库
                Set dsBeanNames = this.resolvedDataSources.keySet();
                for (Object dsBeanName : dsBeanNames) {
                    if (dsBeanName.toString().contains(groupName)) {
                        currentLookupKey = dsBeanName.toString();
                        break;
                    }
                }

            } else {
                //do nothing
                //当currentLookupKey=null,最终会选择默认的数据源(此种场景只有一个组并且只有一个写库)
            }
        }

        if (LOGGER.isDebugEnabled()) {
            LOGGER.debug(">>>>> determine current look up key is {}", currentLookupKey);
        }
        return currentLookupKey;
    }

    @Override
    public Connection getConnection() throws SQLException {
        return getConnection(null, null);
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        return getConnectionInternal(username, password);
    }

    private Connection getConnectionInternal(String username, String password) throws SQLException {
        Connection connection;

        DataSource ds = determineTargetDataSource();

        try {
            if (username == null && password == null) {
                connection = ds.getConnection();
            } else {
                connection = ds.getConnection(username, password);
            }

        } catch (Exception e) {

            if (this.loadBalanceStrategyMap == null) {
                //未开启读写分离不注入负载均衡器,直接抛出

                throw e;

            } else {

                String currentDsName = (String) determineCurrentLookupKey();

                doHeartBeat();

                LoadBalanceStrategy<String> loadBalanceStrategy;
                if (StringUtils.isNotBlank(DbShardingStrategyHolder.getDbGroupName())) {
                    loadBalanceStrategy = this.loadBalanceStrategyMap.get(DbShardingStrategyHolder.getDbGroupName());
                } else {
                    loadBalanceStrategy = this.loadBalanceStrategyMap.values().iterator().next();
                }

                if (StringUtils.equals(currentDsName, loadBalanceStrategy.getWriteDsBeanName())) {
                    // write datasource has problem

                    LOGGER.error("failed to reconnect from master_db [{}]", currentDsName);
                    throw new SQLException(e.getMessage(), e);

                } else {
                    // read datasource has problem

                    loadBalanceStrategy.removeTarget(currentDsName);

                    String readDsName = loadBalanceStrategy.elect();
                    if (StringUtils.isBlank(readDsName)) {
                        readDsName = loadBalanceStrategy.getWriteDsBeanName();
                    }

                    ReadWriteSeparationHolder.setReadWriteDatasourceName(readDsName);

                    return getConnectionInternal(username, password);
                }
            }
        }
        return connection;
    }

    private boolean readWriteEnabled = false;

    public void setReadWriteEnabled(boolean readWriteEnabled) {
        synchronized (this) {
            this.readWriteEnabled = readWriteEnabled;
            if (this.readWriteEnabled) {
                this.heartBeatThreadPool = Executors.newFixedThreadPool(1);
                this.dataSourceHeartBeat = new DataSourceHeartBeat(this);
                this.loadBalanceStrategyMap = this.applicationContext.getBeansOfType(LoadBalanceStrategy.class);

                Assert.notNull(this.loadBalanceStrategyMap, "loadBalanceStrategyMap is required");
            }
        }
    }

    private ExecutorService heartBeatThreadPool;
    private DataSourceHeartBeat dataSourceHeartBeat;

    @Autowired
    private ApplicationContext applicationContext;
    private Map<String, LoadBalanceStrategy> loadBalanceStrategyMap;

    private synchronized void doHeartBeat() {
        if (!this.dataSourceHeartBeat.isRunning()) {
            this.heartBeatThreadPool.execute(this.dataSourceHeartBeat);
        }
    }

    private synchronized void shutdownHeartBeat() {
        this.dataSourceHeartBeat.close();
        this.heartBeatThreadPool.shutdown();
    }

    @PreDestroy
    public void destroy() throws Exception {
        shutdownHeartBeat();
    }

    @Override
    public Set<String> getReadWriteKeys() {
        Set<String> readWriteKeys = new HashSet<>();
        for (Object o : this.resolvedDataSources.keySet()) {
            readWriteKeys.add((String) o);
        }
        return readWriteKeys;
    }

    private static class DataSourceHeartBeat implements Runnable {

        private static final Logger LOGGER = LoggerFactory.getLogger(DataSourceHeartBeat.class);

        private static final int DEFAULT_HEART_BEAT_INTERVAL = 5;

        private RoutingDataSource routingDataSource;

        private volatile boolean running = false;
        private volatile boolean close = false;

        public boolean isRunning() {
            return this.running;
        }

        public void close() {
            this.close = true;
        }

        public DataSourceHeartBeat(RoutingDataSource routingDataSource) {
            this.routingDataSource = routingDataSource;
        }

        @Override
        public void run() {
            this.running = true;

            if (!this.close) {

                Set<LoadBalanceStrategy> lbStrategySet = new LinkedHashSet<>(this.routingDataSource.loadBalanceStrategyMap.values());

                Set<String> failureDsBeanNames = new HashSet<>();
                for (LoadBalanceStrategy<String> lbStrategy : lbStrategySet) {
                    failureDsBeanNames.addAll(lbStrategy.getFailedTargets());

                    //do heart beat to master db every time
                    failureDsBeanNames.add(lbStrategy.getWriteDsBeanName());
                }

                //invoke DataSourceWarningProvider spi impl if exist
                doWarning(failureDsBeanNames);

                boolean recoverOk = true;
                for (; ; ) {
                    for (String dsBeanName : failureDsBeanNames) {
                        LOGGER.info(">>>>>>>>>> Got validate connection [key = " + dsBeanName + "] is ok or not ???");
                        try {
                            DataSource ds = this.routingDataSource.resolvedDataSources.get(dsBeanName);

                            DataSourceUtil.doHeartBeat(ds);

                            for (LoadBalanceStrategy<String> lbStrategy : lbStrategySet) {
                                if (!StringUtils.equals(dsBeanName, lbStrategy.getWriteDsBeanName())) {
                                    lbStrategy.recoverTarget(dsBeanName);
                                }
                            }
                            if (LOGGER.isDebugEnabled()) {
                                LOGGER.debug(">>>>>>>>>> validate connection [key = " + dsBeanName + " ] is ok");
                            }
                        } catch (Exception e) {
                            recoverOk = false;
                            LOGGER.error(">>>>>>>>>> validate connection [key = " + dsBeanName + " ] is failed");
                        }
                    }

                    if (recoverOk) {
                        break;
                    }

                    try {
                        TimeUnit.SECONDS.sleep(DEFAULT_HEART_BEAT_INTERVAL);
                    } catch (Exception e) {
                        LOGGER.error(e.getMessage(), e);
                    }
                }
            }
            this.running = false;
        }

        private void doWarning(Set<String> failureDsNames) {
            List<Map<String, Object>> failureDataSourceInfo = new ArrayList<>();

            for (Map.Entry<Object, DataSource> entry : this.routingDataSource.resolvedDataSources.entrySet()) {
                if (failureDsNames.contains(entry.getKey()) && entry.getValue() instanceof DruidDataSource) {
                    failureDataSourceInfo.add(((DruidDataSource) entry.getValue()).getStatData());
                }
            }

            DataSourceWarningProvider dataSourceWarningProvider = null;

            //java spi mechanism
            ServiceLoader<DataSourceWarningProvider> serviceLoader = ServiceLoader.load(DataSourceWarningProvider.class);
            Iterator<DataSourceWarningProvider> iterator = serviceLoader.iterator();
            while (iterator.hasNext()) {
                dataSourceWarningProvider = iterator.next();
                break;
            }

            if (dataSourceWarningProvider != null) {
                dataSourceWarningProvider.doWarning(failureDataSourceInfo);
            } else {
                Map<String, DataSourceWarningProvider> dataSourceWarningProviderMap = this.routingDataSource.applicationContext.getBeansOfType(DataSourceWarningProvider.class);
                if (dataSourceWarningProvider != null && dataSourceWarningProviderMap.size() > 0) {
                    dataSourceWarningProviderMap.values().iterator().next().doWarning(failureDataSourceInfo);
                }
            }

            if (LOGGER.isDebugEnabled()) {
                LOGGER.debug(">>>>> complete warning ! <<<<<");
            }

        }
    }
}
